package config

import (
	"fmt"
	"log"
	"ota-backend/models"

	"gorm.io/driver/mysql"
	"gorm.io/driver/sqlite"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"
)

var DB *gorm.DB

// InitDatabase 初始化数据库连接
func InitDatabase() error {
	dbType := ViperConfig.GetString("database.type")

	var err error

	switch dbType {
	case "sqlite":
		dbPath := ViperConfig.GetString("database.sqlite.path")
		DB, err = gorm.Open(sqlite.Open(dbPath), &gorm.Config{
			Logger: logger.Default.LogMode(logger.Info),
		})
		if err != nil {
			return fmt.Errorf("failed to connect to sqlite: %v", err)
		}
		log.Println("Connected to SQLite database")

	case "mysql":
		dsn := fmt.Sprintf("%s:%s@tcp(%s:%d)/%s?charset=%s&parseTime=True&loc=Local",
			ViperConfig.GetString("database.mysql.username"),
			ViperConfig.GetString("database.mysql.password"),
			ViperConfig.GetString("database.mysql.host"),
			ViperConfig.GetInt("database.mysql.port"),
			ViperConfig.GetString("database.mysql.database"),
			ViperConfig.GetString("database.mysql.charset"),
		)
		DB, err = gorm.Open(mysql.Open(dsn), &gorm.Config{
			Logger: logger.Default.LogMode(logger.Info),
		})
		if err != nil {
			return fmt.Errorf("failed to connect to mysql: %v", err)
		}
		log.Println("Connected to MySQL database")

	default:
		return fmt.Errorf("unsupported database type: %s", dbType)
	}

	// 自动迁移数据库表
	err = DB.AutoMigrate(
		&models.Product{},
		&models.Firmware{},
		&models.Task{},
		&models.UpgradeRecord{},
		&models.Device{},
		&models.LoginFailure{},
		&models.IPBlacklist{},
	)
	if err != nil {
		return fmt.Errorf("failed to migrate database: %v", err)
	}

	log.Println("Database migration completed")
	return nil
}
